"""
提供生成论文的API接口

简单模式，一次只执行一个任务
"""
import os
import sys
import glob
import time
import signal
import sqlite3
import logging
import argparse
import multiprocessing

from datetime import datetime
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from contextlib import asynccontextmanager
from pydantic import BaseModel
from typing import List
import uvicorn
from utils import AIWriter, Convert2pdf


class StartRequest(BaseModel):
    ids: List[int]
    model: str = "gpt-oss:20b"

def dict_factory(cursor, row):
    d = {}
    for idx, col in enumerate(cursor.description):
        d[col[0]] = row[idx]
    return d


def write_task(logger, model, records, output_dir):
    try:
        start_time = datetime.now()
        # 清理目录下的MD、Tex、PDF文件
        logger.info("Clean MD files...")
        for file_path in glob.glob(os.path.join(output_dir, "*.md")):
            os.remove(file_path)

        logger.info("Clean Tex files...")
        for file_path in glob.glob(os.path.join(output_dir, "*.tex")):
            os.remove(file_path)

        logger.info("Clean PDF files...")
        for file_path in glob.glob(os.path.join(output_dir, "*.pdf")):
            os.remove(file_path)

        logger.info("Starting to generate paper...")
        ai_writer = AIWriter(
            logger=logger,
            model=model,
            records=records,
            output_dir=output_dir,
        )
        ai_writer.auto_generate()
        
        logger.info("Starting to create the PDF file...")
        Convert2pdf(
            temp_dir=output_dir,
            title=ai_writer.title,
            keywords=ai_writer.keywords,
        )
    except Exception as e:
        logger.warning(str(e))
    finally:
        logger.info(f"Total execution time: {(datetime.now() - start_time) / 60} m")


def test_process():
    time.sleep(10)

@asynccontextmanager
async def lifespan(app: FastAPI):
    # 应用启动时执行
    app.state.process = None
    app.state.agent_path = "/data/agents/paper_agent"
    os.makedirs(app.state.agent_path, exist_ok=True)

    log_file = os.path.join(app.state.agent_path, "run.log")
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
        handlers=[
            logging.FileHandler(log_file, encoding="utf-8", mode="w"),
            logging.StreamHandler(sys.stdout),
        ],
    )

    app.state.logging = logging

    # 初始化agent目录
    os.makedirs(app.state.agent_path, exist_ok=True)

    yield
    # 应用关闭时执行
    if app.state.process and app.state.process.is_alive():
        pid = app.state.process.pid
        os.kill(pid, signal.SIGKILL)



app = FastAPI(lifespan=lifespan)

@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    return JSONResponse(
        status_code=500,
        content={
            "status": "error",
            "code": 500,
            "message": "Internal Server Error",
            "details": str(exc)
        },
    )


@app.post("/start")
def start(req: StartRequest):
    haiant_conn = None
    try:
        # 查询数据
        db_path = "/data/personal_ai_app.db"
        if os.path.exists(db_path):
            haiant_conn = sqlite3.connect(db_path, check_same_thread=False)
            haiant_conn.row_factory = dict_factory
            haiant_cursor = haiant_conn.cursor()

            # 根据ID查询多条记录
            sql = "SELECT * FROM Record WHERE id IN ({})".format(", ".join(["?"] * len(req.ids)))
            haiant_cursor.execute(sql, req.ids)
            records = haiant_cursor.fetchall()
            if haiant_conn: haiant_conn.close()

            if len(records) > 0:
                process = multiprocessing.Process(target=write_task, args=(app.state.logging, req.model, records, app.state.agent_path))
                # process = multiprocessing.Process(target=test_process)
                process.start()
                app.state.process = process
                return {"status": "success", "message": "start to write paper"}
            else:
                return {"status": "error", "message": "can't find records"}

    except Exception as e:
        app.state.logging.warning(str(e))
        return {"status": "error", "message": str(e)}
    finally:
        if haiant_conn:
            haiant_conn.close()

@app.post("/stop")
def stop():
    try:
        # # 检查进程是否存在
        if app.state.process and app.state.process.is_alive():
            pid = app.state.process.pid
            os.kill(pid, signal.SIGKILL)
            app.state.process = None
            app.state.logging.info(f"The process with PID {pid} has been stopped.")
            return {"status": "success", "message": "Task already stopped."}
        else:
            return {"status": "success", "message": "No running task"}
    except Exception as e:
        app.state.logging.error(f"Error stopping task: {str(e)}")
        return {"status": "error", "message": f"Error stopping task: {str(e)}"}


@app.get("/status")
def status():
    try:
        # 检查进程执行状态
        if app.state.process and app.state.process.is_alive():
            return {"status": "success", "data": {"is_alive": True}}
        else:
            return {"status": "success", "data": {"is_alive": False}}
    except Exception as e:
        return {"status": "error", "message": f"Error stopping task: {str(e)}"}


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--port", type=int, default=8080, help="端口号，默认8080")
    args = parser.parse_args()
    
    # start the FastAPI app
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=args.port
    )

    print(f"Server is running on http://localhost:{args.port}")
